package com.wuwangfu.state;

import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.state.StateTtlConfig;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

/**
 * @Author: jcshen
 * @Date: 2023-03-09
 * @PackageName: com.wuwangfu.state
 * @ClassName: KeyedStateTTL
 * @Description: 设置 ValueState 的存活时间
 * @Version: 1.0.0
 *
 * https://nightlies.apache.org/flink/flink-docs-release-1.14/docs/dev/datastream/fault-tolerance/state/#using-keyed-state
 *
 */
public class KeyedStateTTL {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.enableCheckpointing(10000);
        env.setRestartStrategy(RestartStrategies.fixedDelayRestart(3,5000));

        DataStreamSource<String> lines = env.socketTextStream("localhost", 8888);

        SingleOutputStreamOperator<Tuple2<String, Integer>> flated = lines.flatMap(new FlatMapFunction<String, Tuple2<String, Integer>>() {
            @Override
            public void flatMap(String value, Collector<Tuple2<String, Integer>> out) throws Exception {
                String[] fields = value.split(" ");
                for (String word : fields) {
                    if ("error".equals(word)) {
                        throw new RuntimeException("出错了");
                    }
                    out.collect(Tuple2.of(word, 1));
                }
            }
        });
        //
        KeyedStream<Tuple2<String, Integer>, String> keyed = flated.keyBy(t -> t.f0);
        //
        keyed.map(new RichMapFunction<Tuple2<String, Integer>, Tuple2<String, Integer>>() {

            private transient ValueState<Integer> valueState;

            @Override
            public void open(Configuration parameters) throws Exception {
                ValueStateDescriptor<Integer> valueDesc = new ValueStateDescriptor<>("wc-state", Integer.class);
                //定义一个状态 TtlConfig
                StateTtlConfig ttlConfig = StateTtlConfig.newBuilder(Time.seconds(10))
                        //开始计时
                        .setUpdateType(StateTtlConfig.UpdateType.OnCreateAndWrite)
                        //超时不可见
                        .setStateVisibility(StateTtlConfig.StateVisibility.NeverReturnExpired)
                        .build();
                //关联状态描述器
                valueDesc.enableTimeToLive(ttlConfig);

                valueState = getRuntimeContext().getState(valueDesc);
            }

            @Override
            public Tuple2<String, Integer> map(Tuple2<String, Integer> tuple) throws Exception {
                Integer state = valueState.value();
                if (state == null){
                    state = 0;
                }
                int total = tuple.f1 + state;
                //更新状态
                valueState.update(total);
                //输出
                tuple.f1 = total;
                return tuple;
            }
        }).print();

        env.execute();
    }
}
